import streamlit as st

import dns.resolver

import pandas as pd

from io import BytesIO

import time

from collections import defaultdict

# 设置页面标题和图标

st.set_page_config(

    page_title="企业邮箱品牌识别工具",

    page_icon="📧",

    layout="wide"

)

# 扩展的企业邮箱品牌特征映射

EMAIL_PROVIDERS = {
    '腾讯qq邮箱':['mx1.qq.com','mx2.qq.com','mx3.qq.com'],
    '阿里企业邮箱': ['mxn.mxhichina.com', 'mxw.mxhichina.com', 'mx1.qiye.aliyun.com', 'mx2.qiye.aliyun.com', 'mx3.qiye.aliyun.com'],
    '钉钉企业邮箱':['em.dingtalk.com'],
    '阿里邮箱国产化版':['xc.aliyun.com'],
    '阿里巴巴外贸邮':['mail.alibaba.com'],
    '腾讯企业邮箱': ['mxbiz1.qq.com', 'mxbiz2.qq.com'],
    '网易企业邮箱': ['qiye163mx01.mxmail.netease.com', 'qiye163mx02.mxmail.netease.com','mxmail.netease.com'],
    '网易免费企业邮': ['mx.ym.163.com','163.com'],
    '新浪企业邮箱': ['mx.sina.net', 'mx.entmail.sina.com', 'mx.sinanet.com'],
    '新网企业邮箱':['mx1.xinnet.com','mx2.xinnet.com','global-mail.cn'],
    '飞书企业邮箱':['mx1.feishu.cn','mx2.feishu.cn','mx3.feishu.cn'],
    'Coremail企业邮箱':['mx.c3.icoremail.net','mx.c2.icoremail.net','mx.c1.icoremail.net','mx.corpease.net','mx-coremail-cn.icoremail.net'],
    '263企业邮箱':['mxwcom.263xmail.com','mxcom.263xmail.com','mx.263.net','smtp.263xmail.com','smtp.263mail.com'],
    '35互联企业邮箱':['qiye.mail.35.com','mx1.yunyou.top','mx2.yunyou.top','mx3.yunyou.top'],
    'Zoho邮箱': ['mx.zoho.com.cn', 'mx2.zoho.com.cn','mx1.zoho.com','mx2.zoho.com','mx3.zoho.com'],
    '易名企业邮箱':['mx1.enom.com','mx2.enom.com'],
    '华为云企业邮箱': ['mx1.sparkspace.huaweicloud.com', 'mx2.sparkspace.huaweicloud.com'],
    '谷歌企业邮箱（Gmail）': ['google.com', 'googlemail.com', 'aspmx.l.google.com'],
    '微软Office 365': ['outlook.com', 'protection.outlook.com', 'mail.protection.outlook.com','outlook.cn'],
    '中国电信企业邮箱（21CN）': ['mta-ent.21cn.com','mailcdn.21cn.com'],
    'TOM企业邮箱':['mx.qy.tom.com'],
    '亚马逊WorkMail': ['amazonaws.com', 'awsapps.com'],
    'FastMail': ['fastmail.com', 'messagingengine.com'],
    'Yandex Mail': ['yandex.com', 'yandex.ru', 'yandex.net'],
    'Mail.ru': ['mail.ru', 'corp.mail.ru'],
    'GoDaddy企业邮箱':['mx.secureserver.net'],
    '晶格云企业邮箱':['ee-post.com'],
    '老板邮局-中资源':['chinaemail.cn'],
    '自建邮件服务器': []  # 默认类别

}

# 添加一些常见的邮件服务器关键字到自建类别

SELF_HOSTED_KEYWORDS = ['mail', 'smtp', 'mx', 'exchange', 'postfix', 'dovecot', 'zimbra', 'coremail']


def check_mx_records(domain):
    """查询域名的MX记录"""

    try:

        mx_records = dns.resolver.resolve(domain, 'MX')

        return [str(mx.exchange).rstrip('.').lower() for mx in mx_records]

    except (dns.resolver.NoAnswer, dns.resolver.NXDOMAIN,

            dns.resolver.NoNameservers, dns.resolver.Timeout,

            dns.exception.DNSException):

        return []


def identify_email_provider(domain):
    """识别企业邮箱品牌"""

    mx_records = check_mx_records(domain)

    if not mx_records:
        return "无法获取MX记录或没有设置MX记录", mx_records

    # 检查每个MX记录匹配的品牌

    matched_providers = set()

    for provider, keywords in EMAIL_PROVIDERS.items():

        for mx in mx_records:

            if any(keyword in mx for keyword in keywords):
                matched_providers.add(provider)

    # 检查是否是自建邮件服务器

    is_self_hosted = False

    for mx in mx_records:

        if any(keyword in mx for keyword in SELF_HOSTED_KEYWORDS):

            # 但需要排除已知服务商的关键字

            known_provider_found = False

            for provider, keywords in EMAIL_PROVIDERS.items():

                if provider != '自建邮件服务器' and any(keyword in mx for keyword in keywords):
                    known_provider_found = True

                    break

            if not known_provider_found:
                is_self_hosted = True

                break

    if matched_providers:

        return ", ".join(matched_providers), mx_records

    elif is_self_hosted:

        return "自建邮件服务器", mx_records

    else:

        return f"未知服务商", mx_records


def process_domains(domains):
    """处理域名列表并返回结果"""

    results = []

    progress_bar = st.progress(0)

    status_text = st.empty()

    for i, domain in enumerate(domains):

        domain = domain.strip()

        if not domain:
            continue

        status_text.text(f"正在查询: {domain} ({i + 1}/{len(domains)})")

        provider, mx_records = identify_email_provider(domain)

        results.append({

            "域名": domain,

            "邮箱品牌": provider,

            "MX记录": ", ".join(mx_records) if mx_records else "无"

        })

        # 更新进度条

        progress_bar.progress((i + 1) / len(domains))

        # 添加短暂延迟避免DNS查询过于频繁

        time.sleep(0.1)

    status_text.text("查询完成!")

    return results


def main():
    st.title("📧 企业邮箱品牌识别工具")

    st.markdown("""

    本工具可以帮助您批量识别域名使用的企业邮箱品牌。只需输入域名列表，工具会自动查询MX记录并识别邮箱服务商。

    """)

    # 创建两列布局

    col1, col2 = st.columns([1, 1])

    with col1:

        st.header("输入域名")

        input_method = st.radio(

            "选择输入方式",

            ["文本输入", "文件上传"]

        )

        domains = []

        if input_method == "文本输入":

            domain_input = st.text_area(

                "请输入域名列表（每行一个域名）",

                height=200,

                placeholder="example.com\ngoogle.com\nmicrosoft.com"

            )

            if domain_input:
                domains = [line.strip() for line in domain_input.split('\n') if line.strip()]

        else:

            uploaded_file = st.file_uploader("上传文本文件", type=['txt'])

            if uploaded_file is not None:
                domains = [line.decode('utf-8').strip() for line in uploaded_file.readlines() if line.strip()]

        if st.button("开始查询", type="primary") and domains:

            with st.spinner("正在查询，请稍候..."):

                results = process_domains(domains)

                if results:
                    # 显示结果表格

                    st.header("查询结果")

                    df = pd.DataFrame(results)

                    st.dataframe(df, use_container_width=True)

                    # 添加下载按钮

                    csv = df.to_csv(index=False).encode('utf-8')

                    st.download_button(

                        label="下载结果为CSV",

                        data=csv,

                        file_name="email_provider_results.csv",

                        mime="text/csv",

                    )

                    # 显示统计信息

                    st.subheader("统计信息")

                    provider_counts = df['邮箱品牌'].value_counts()

                    st.dataframe(provider_counts.rename_axis('邮箱品牌').reset_index(name='数量'))

                    # 显示图表

                    st.bar_chart(provider_counts)

    with col2:

        st.header("使用提示")

        st.markdown("""

        1. 输入域名时不需要添加"http://"或"www"前缀

        2. 批量查询时，建议不要一次性处理过多域名（建议不超过100个）

        3. 查询结果基于域名的MX记录，某些企业可能使用自定义MX记录

        4. 对于无法识别的服务商，工具会显示MX记录供参考

        5. 查询结果可以导出为CSV文件进行进一步分析

        """)


if __name__ == "__main__":
    main()